import numpy as np
from bert4keras.backend import keras,set_gelu,K
from bert4keras.tokenizers import Tokenizer
from bert4keras.models import build_transformer_model
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding,DataGenerator
from bert4keras.snippets import open
from keras.layers import Dropout,Dense

set_gelu('tanh')

maxlen = 128
batch_size = 64
import os
import sys
if 'win' in sys.platform:
    BertPath = r'E:\nlp-data\bert\model_chinese'
else:
    BertPath = r'/data/nlp_data/bert/model_chinese/'

config_path = os.path.join(BertPath,'bert_config.json')
checkpoint_path = os.path.join(BertPath,'bert_model.ckpt')
dict_path = os.path.join(BertPath,'vocab.txt')

def load_data(filename):
    D = []
    with open(filename,encoding='utf-8') as f:
        for l in f:
            text1,text2,label = l.strip().split("\t")
            D.append((text1,text2,int(label)))

    return D

train_data = load_data('datasets/lcqmc/train.txt')
valid_data = load_data('datasets/lcqmc/dev.txt')
test_data = load_data('datasets/lcqmc/test.txt')

tokenizer = Tokenizer(dict_path)

class data_generator(DataGenerator):
    def __iter__(self,random=False):
        batch_token_ids,batch_segment_ids,batch_labels = [],[],[]
        for is_end,(text1,text2,label) in self.sample(random):
            token_ids,segment_ids = tokenizer.encode(
                text1,text2,max_length=maxlen
            )

            batch_token_ids.append(token_ids)
            batch_segment_ids.append(segment_ids)
            batch_labels.append([label])

            if len(batch_token_ids) == self.batch_size or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids =sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids,batch_segment_ids],batch_labels
                batch_token_ids,batch_segment_ids,batch_labels = [],[],[]

bert = build_transformer_model(
    config_path=config_path,
    checkpoint_path=checkpoint_path,
    with_pool=True,
    return_keras_model=False
)

output = Dropout(rate=0.1)(bert.model.output)
output = Dense(
    units=2,activation='softmax',kernel_initializer=bert.initializer
)(output)

model = keras.models.Model(bert.model.input,output)
model.summary()

model.compile(
    loss='sparse_categorical_crossentropy',
    optimizer=Adam(2e-5),
    metrics=['accuracy']
)

# 转换数据集
train_generator = data_generator(train_data, batch_size)
valid_generator = data_generator(valid_data, batch_size)
test_generator = data_generator(test_data, batch_size)

def evaluate(data):
    total, right = 0., 0.
    for x_true, y_true in data:
        y_pred = model.predict(x_true).argmax(axis=1)
        y_true = y_true[:, 0]
        total += len(y_true)
        right += (y_true == y_pred).sum()
    return right / total

class Evaluator(keras.callbacks.Callback):
    def __init__(self):
        self.best_val_acc = 0

    def on_epoch_end(self, epoch, logs=None):
        val_acc = evaluate(valid_generator)
        if val_acc > self.best_val_acc:
            self.best_val_acc = val_acc
            model.save_weights('output/lcqmc.weights')
        test_acc = evaluate(test_generator)
        print(
            u'val_acc: %.5f, best_val_acc: %.5f, test_acc: %.5f\n' %
            (val_acc, self.best_val_acc, test_acc)
        )

if __name__ == '__main__':
    if not os.path.exists('output/lcqmc.weights'):
        evaluator = Evaluator()
        model.fit_generator(
            train_generator.forfit(),
            steps_per_epoch=len(train_generator),
            epochs=20,
            callbacks=[evaluator]
        )
    else:
        train_generator = data_generator(train_data, batch_size)
        for i in train_generator.forfit():
            print(i)
            break
        model.load_weights('output/lcqmc.weights')
        print(u'final test acc: %05f\n' % (evaluate(test_generator)))